# File: arima.R
# R source code for fitting/diagnosing ARMA models.
# Author: Congxing Cai (congxing@stanford.edu) 
ts.acf.plots <- function(ts, filename) {
	# Plots the time series, and its acf & pacf values.
	# Args:
	#   ts: the time series object to be plotted.
	#   filename: the png file to save the plots to.
	png(filename)
	par(mfrow=c(3,1))
	ts.plot(ts)
	acf(ts)
	pacf(ts)
	dev.off()
}

ts.diff.plots <- function(ts, d, filename) {
	# Plots the time series after differencing (till d-level)
	# Args:
	#   ts: the time series object to be differenced.
	#   d: the maximum degree for differencing.
	#   filename: the png file to save the plots to.
	png(filename)
	par(mfrow=c(d,1))
	for (i in 1:d) {
		ts <- diff(ts)
		ts.plot(ts)
	}
	dev.off()
}

ts.stl.plots <- function(ts, filename) {
	# Plot the seasonal decompositions of time series.
	# Args:
	#   ts: the time series object.
	#   filename: the png file to save the plots to.
	# Returns:
	#   a list with 3 elements: seasonal, trend, remainder (series)
	png(filename)
	ts.stl <- stl(ts, "periodic")
	par(mfrow=c(3,1))
	ts.plot(ts.stl$time[,1], gpars=list(ylab="seasonal"))
	ts.plot(ts.stl$time[,2], gpars=list(ylab="trend"))
	ts.plot(ts.stl$time[,3], gpars=list(ylab="remainder"))
	dev.off()
	return(list(seasonal=ts.stl$time[,1], trend=ts.stl$time[,2], remainder=ts.stl$time[,3]))
}

arma.model.selection <- function(ts, pmax, qmax, d=0) {
	# Find the best model parameter (p, d, q) for arima based on the minimum aic
	# Args:
	#   ts: the time series object.
	#   pmax: the max value of p to search.
	#   qmax: the max value of q to search.
	#   d: the d value for integrated ARMA, default is 0
	# Returns:
	#   a list best model parameter: (p=p_minaic, q=q_minaic)
	aic<-matrix(rep(0,pmax*qmax), pmax, qmax);
	for (i in 1:pmax) for (j in 1:qmax) {
		fit.arima <- arima(ts, order=c(i, d, j), method="ML")
		aic[i,j]<-fit.arima$aic
	}
	#print(aic)
	minaic.ind <- which(aic==min(aic), arr.ind=TRUE)
	return(list(p=minaic.ind[1,"row"], q=minaic.ind[1,"col"]))
}

arma.model.diag <- function(arma.model, filename) {
	# Plot the diagnose charts into the file for the arma.model.
	# Args:
	#   arma.model: the fitted arma model.
	#   filename: the png file to save charts.
	png(filename)
	tsdiag(arma.model)
	dev.off()	
}

model.eval <- function(ts.pred, ts.se, ts.truth, filename) {
	# Evaluate the predicted time series with the truth.
	# Args:
	#   ts.pred: time series of the predicted values.
	#   ts.se: time series of the standard errors.
	#   ts.truth: time series of the true values.
	#   filename: the png file to save the charts
	# Retrun:
	#   a list of measures: (correct, wrong, precision)
	png(filename)
	ts.plot(ts.pred, ts.pred+ts.se, ts.pred-ts.se, gpars=list(lty=c(1,2,2)))
	points(ts.truth, col="red")
	dev.off()
	correct = sum((ts.truth <= ts.pred + ts.se) & (ts.truth >= ts.pred - ts.se))
	wrong = length(ts.truth) - correct
	precision = correct / length(ts.truth)
	return(list(correct=correct, wrong=wrong, precision=precision))
}
